# Description: The code is to find the optimal path of the codon lattice.
# Author: Shibo Li, MiQroEra
# Date: 2023-07-15
# PS: The code is based on the dynamic programming algorithm.

import numpy as np
from utils.evaluate import calculate_mfe
def find_optimal_path(codon_lattice, codon_usage):
    num_amino_acids = len(codon_lattice)
    num_codons = len(codon_lattice[0])
    dp = np.zeros((num_amino_acids, num_codons))
    prev = np.zeros((num_amino_acids, num_codons), dtype=int)

    for i in range(num_amino_acids):
        for j in range(num_codons):
            codon = codon_lattice[i][j]
            cai = codon_usage.get(codon, 0)
            mfe = calculate_mfe(codon)
            score = cai * mfe  # 综合评价函数，综合考虑CAI和MFE
            
            if i == 0:
                dp[i][j] = score
            else:
                max_score = -float('inf')
                max_idx = -1
                for k in range(num_codons):
                    prev_score = dp[i - 1][k]
                    transition_score = cai * calculate_mfe(codon_lattice[i - 1][k] + codon)  # 过渡得分也考虑CAI和MFE
                    curr_score = prev_score + transition_score
                    if curr_score > max_score:
                        max_score = curr_score
                        max_idx = k
                dp[i][j] = max_score
                prev[i][j] = max_idx
    
    # 回溯最优路径
    path = []
    idx = np.argmax(dp[-1])
    for i in range(num_amino_acids - 1, -1, -1):
        path.append(codon_lattice[i][idx])
        idx = prev[i][idx]
    
    # 将路径反转并拼接成mRNA序列
    mrna_sequence = ''.join(reversed(path))
    
    return mrna_sequence
